package net.lovesrc.allin.biz.lbc;

import net.coobird.thumbnailator.Thumbnails;
import net.sourceforge.tess4j.ITessAPI;
import net.sourceforge.tess4j.ITesseract;
import net.sourceforge.tess4j.Tesseract;
import net.sourceforge.tess4j.TesseractException;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author clg
 * Created at 21:23 2019/3/4
 * @since 1.1
 */
public class Captcha {

    private String dataPath;

    public Captcha(String dataPath) {
        this.dataPath = dataPath;
    }

    public synchronized String process(InputStream is) throws IOException, TesseractException {
        //size 120*36
        //split
        List<BufferedImage> images = split(is, 4);

        List<String> targetResults = Arrays.asList("0", "1", "2", "3", "4", "5", "6", "7", "8", "9");

        StringBuffer sb = new StringBuffer();
        //识别,如果不是0-9重复识别
        for (BufferedImage bufferedImage : images) {
            String result = recognition(targetResults, bufferedImage);
            sb.append(result);
        }

        //输出结果
        return sb.toString();
    }

    private List<BufferedImage> split(InputStream is, int size) throws IOException {
        BufferedImage bufferedImage = ImageIO.read(is);
        if (bufferedImage == null) {
            return new ArrayList<>();
        }
        int width = bufferedImage.getWidth() - 5;
        int height = bufferedImage.getHeight();
        List<BufferedImage> images = new ArrayList<>();
        for (int i = 0; i < size; i++) {
            int x = (width / size) * i;
            BufferedImage image = Thumbnails.of(bufferedImage)
                    .sourceRegion(x, 0, width / size, height)
                    .size(width / size, height)
                    .outputQuality(1.0)
                    .asBufferedImage();
            images.add(image);
        }
        return images;
    }

    private ITesseract tesseract;

    private ITesseract getTesseract() {
        if (tesseract == null) {
            tesseract = new Tesseract();
            tesseract.setDatapath(dataPath);
            tesseract.setPageSegMode(10);
            tesseract.setOcrEngineMode(ITessAPI.TessOcrEngineMode.OEM_TESSERACT_ONLY);
            tesseract.setTessVariable("tessedit_char_whitelist", "0123456789");
        }
        return tesseract;
    }

    private String recognition(List<String> targetResults, BufferedImage BufferedImage) throws TesseractException {
        for (int i = 100; i < 355; i = i + 25) {
            java.awt.image.BufferedImage image = new OcrFilter(BufferedImage).medianX().medianY().extrem(i % 255).optBufferedImage();
            String result = getTesseract().doOCR(image);
            if (result == null) {
                continue;
            }
            result = result.trim();
            if (targetResults.contains(result)) {
                return result;
            }
        }
        return targetResults.get(0);
    }
}
